import asyncio
import os
from abc import ABC, abstractmethod
from typing import AsyncGenerator, Any, Optional
from routes.local.streaming.stream_config import OutputType, MediaType

from utils.logger import setup_logger
logger = setup_logger()

from utils.file_handler import write_file_stream, get_mime_type, get_file_type


class OutputHandler(ABC):
    @abstractmethod
    async def process(self, data: AsyncGenerator[Any, None], output_type: OutputType) -> AsyncGenerator[bytes, None]:
        pass

class TextStreamOutputHandler(OutputHandler):
    async def process(self, data: AsyncGenerator[Any, None], output_type: OutputType) -> AsyncGenerator[bytes, None]:
        logger.info("Generating text stream output")
        async for item in data:
            if isinstance(item, bytes):
                try:
                    # 尝试解码为文本
                    text = item.decode('utf-8')
                    yield text.encode('utf-8')
                except UnicodeDecodeError:
                    # 如果解码失败，保持原样
                    yield item
            else:
                yield str(item).encode('utf-8')
            await asyncio.sleep(0.01)

class AudioStreamOutputHandler(OutputHandler):
    async def process(self, data: AsyncGenerator[Any, None], output_type: OutputType) -> AsyncGenerator[bytes, None]:
        logger.info("Generating audio stream output")
        async for chunk in data:
            # 确保音频数据是二进制格式
            if isinstance(chunk, str):
                try:
                    # 尝试将Base64编码的字符串转换为二进制
                    import base64
                    chunk = base64.b64decode(chunk)
                except Exception:
                    # 如果不是Base64，则简单编码
                    chunk = chunk.encode('utf-8')
            yield chunk
            await asyncio.sleep(0.01)

class VideoStreamOutputHandler(OutputHandler):
    async def process(self, data: AsyncGenerator[Any, None], output_type: OutputType) -> AsyncGenerator[bytes, None]:
        logger.info("Generating video stream output")
        async for chunk in data:
            # 确保视频数据是二进制格式
            if isinstance(chunk, str):
                try:
                    # 尝试将Base64编码的字符串转换为二进制
                    import base64
                    chunk = base64.b64decode(chunk)
                except Exception:
                    # 如果不是Base64，则简单编码
                    chunk = chunk.encode('utf-8')
            yield chunk
            await asyncio.sleep(0.01)

class FileOutputHandler(OutputHandler):
    async def process(self, data: AsyncGenerator[Any, None], output_type: OutputType) -> AsyncGenerator[bytes, None]:
        logger.info(f"Generating {output_type} file output")
        # 根据输出类型确定文件扩展名
        ext = "bin"  # 默认二进制文件
        if output_type == OutputType.TEXT_FILE:
            ext = "txt"
        elif output_type == OutputType.AUDIO_FILE:
            ext = "mp3"
        elif output_type == OutputType.VIDEO_FILE:
            ext = "mp4"
        elif output_type == OutputType.IMAGE_FILE:
            ext = "jpg"
            
        file_path = f"output_{id(data)}.{ext}"
        async for item in write_file_stream(file_path, data):
            yield item

def get_output_handler(output_type: OutputType) -> OutputHandler:
    handlers = {
        OutputType.TEXT_STREAM: TextStreamOutputHandler(),
        OutputType.AUDIO_STREAM: AudioStreamOutputHandler(),
        OutputType.VIDEO_STREAM: VideoStreamOutputHandler(),
        OutputType.TEXT_FILE: FileOutputHandler(),
        OutputType.AUDIO_FILE: FileOutputHandler(),
        OutputType.VIDEO_FILE: FileOutputHandler(),
        OutputType.IMAGE_FILE: FileOutputHandler(),
        OutputType.FILE: FileOutputHandler(),
    }
    return handlers.get(output_type, TextStreamOutputHandler())